from mypy import argv,argc,tryx,now,load,save,void

config_o = load('../tmp/config.json')
assert config_o, 'Please check ../tmp/config.json'
#print(config_o)

root_data_path_1min = config_o[f'root_data_path_1min']
#root_data_path_tick = config_o[f'root_data_path_tick']

data_source = 'tdx'

code = argv[1]

import duckdb
con = duckdb.connect(database=':memory:')

def rawv(code):
  codex = code[:-1] + '[' + code[-1:] + ']'
  coden = code.replace('.','_')
  tryx(lambda:con.execute("create view v_{} as select close as p,volume as v,tt as t from '{}/{}*/{}.parquet' WHERE P>0 ORDER BY T".format(coden, root_data_path_1min, data_source, codex)),void)
  return con

def apv(code, part_where='1=1', part_select='*'):
  tbn = 'v_{}'.format(code.replace('.','_'))
  sql = 'SELECT {} FROM {} WHERE {}'.format(part_select, tbn, part_where)
  #print(sql)
  return rawv(code).execute(sql)

import numpy as np
import pandas as pd

def get_tbz(code):
  return apv(code,'1=1','t,p, log2(p / lag(p,1,p) over (rows between 1 preceding and 0 following)) as z')

def get_zt_a(code):
  tbz=get_tbz(code)
  zt_a0 = tbz.fetchnumpy()
  return np.vstack((zt_a0['z'],zt_a0['t'])).T

def get_df(code):
  df = get_tbz(code).df()
  df['time'] = pd.DatetimeIndex((df.t+8*3600000)*1000000)
  df['nrm']=df['z'].cumsum().apply(lambda x:2**x)     # normalized
  return df

n = get_zt_a(code)
print(n)

df = get_df(code)
print(df)

import matplotlib.pyplot as plt

plt.plot(df.time, df.nrm)
plt.show()

plt.plot(df.time, df.p)
plt.show()
